import os
import numpy as np
from osgeo import gdal
from tqdm import tqdm
from scipy import stats
import joblib
from sklearn.linear_model import LinearRegression
# 

# -------------------- 工具函数 --------------------

def read_data_from_directory(directory, keyword=None, extension=".tif"):
    """
    从指定目录读取数据文件，并自动获取NoData值替换为NaN
    :param directory: 数据文件所在的目录
    :param keyword: 文件名中的关键字（可选，用于筛选特定文件）
    :param extension: 文件扩展名（默认读取.tif文件）
    :return: 数据列表（每个元素是一个NumPy数组）
    """
    data_list = []
    file_paths = []

    # 遍历目录下的所有文件
    for file_name in os.listdir(directory):
        # 检查文件扩展名和关键字
        if file_name.endswith(extension) and (keyword is None or keyword in file_name):
            file_path = os.path.join(directory, file_name)
            file_paths.append(file_path)

    # 按文件名排序，确保数据按时间顺序排列
    file_paths.sort()

    # 读取数据
    for file_path in file_paths:
        dataset = gdal.Open(file_path)
        if dataset is None:
            raise FileNotFoundError(f"文件 {file_path} 无法读取")

        band = dataset.GetRasterBand(1)
        data = band.ReadAsArray()

        # 获取NoData值
        no_data_value = band.GetNoDataValue()

        # 如果存在NoData值，将其替换为NaN
        if no_data_value is not None:
            data[data == no_data_value] = np.nan

        # 检查是否存在inf值，并将其替换为NaN
        if np.isinf(data).any():
            data[np.isinf(data)] = np.nan

        data_list.append(data)

    return data_list


def get_template_file(directory, extension=".tif"):
    """
    获取指定目录下的第一个.tif文件作为模板文件
    :param directory: 目录路径
    :param extension: 文件扩展名（默认为.tif）
    :return: 模板文件路径
    """
    # 获取目录下所有.tif文件
    tif_files = [f for f in os.listdir(directory) if f.endswith(extension)]

    if not tif_files:
        raise FileNotFoundError(f"目录 {directory} 中没有找到{extension}文件")

    # 返回第一个.tif文件的完整路径
    return os.path.join(directory, tif_files[0])


def save_image(data, file_name, template_file, output_dir):
    """
    保存图像数据，并使用模板文件的NoData值
    :param data: 要保存的数据（NumPy数组）
    :param file_name: 保存的文件名
    :param template_file: 模板文件路径（用于获取NoData值和地理信息）
    :param output_dir: 输出目录
    """
    # 创建输出目录（如果不存在）
    os.makedirs(output_dir, exist_ok=True)

    # 构建完整的文件路径
    file_path = os.path.join(output_dir, file_name)

    # 打开模板文件
    template_dataset = gdal.Open(template_file)
    if template_dataset is None:
        raise FileNotFoundError(f"模板文件 {template_file} 未找到")

    # 获取模板文件的NoData值
    band = template_dataset.GetRasterBand(1)
    no_data_value = band.GetNoDataValue()

    # 创建输出文件
    driver = gdal.GetDriverByName('GTiff')
    dataset = driver.Create(file_path, data.shape[1], data.shape[0], 1, gdal.GDT_Float32)

    # 设置地理信息和投影
    dataset.SetGeoTransform(template_dataset.GetGeoTransform())
    dataset.SetProjection(template_dataset.GetProjection())

    # 写入数据
    band = dataset.GetRasterBand(1)
    data[np.isnan(data)] = no_data_value  # 将NaN替换为NoData值
    band.WriteArray(data)

    # 设置NoData值
    band.SetNoDataValue(no_data_value)
    dataset.FlushCache()


def clip_ndvi(ndvi_data):
    """
    将NDVI值裁剪到[0, 1]范围内
    :param ndvi_data: 输入的NDVI数据（NumPy数组）
    :return: 裁剪后的NDVI数据（NumPy数组）
    """
    # 将大于1的值设置为1
    ndvi_data[ndvi_data > 1] = 1
    # 将小于0的值设置为0
    ndvi_data[ndvi_data < 0] = 0
    return ndvi_data


def save_model(model, coefficients, intercepts, r2_values, t_values, p_values, output_dir="model"):
    """
    保存模型及其检验值到指定目录
    :param model: 回归模型
    :param coefficients: 回归系数（形状为(rows, cols, 3)）
    :param intercepts: 截距（形状为(rows, cols)）
    :param r2_values: R²值（形状为(rows, cols)）
    :param t_values: t值（形状为(rows, cols, 4)）
    :param p_values: p值（形状为(rows, cols, 4)）
    :param output_dir: 输出目录（默认为"model"）
    """
    # 创建输出目录（如果不存在）
    os.makedirs(output_dir, exist_ok=True)

    # 保存模型
    model_path = os.path.join(output_dir, "regression_model.pkl")
    joblib.dump(model, model_path)

    # 保存回归系数
    np.save(os.path.join(output_dir, "coefficients.npy"), coefficients)

    # 保存截距
    np.save(os.path.join(output_dir, "intercepts.npy"), intercepts)

    # 保存R²值
    np.save(os.path.join(output_dir, "r2_values.npy"), r2_values)

    # 保存t值
    np.save(os.path.join(output_dir, "t_values.npy"), t_values)

    # 保存p值
    np.save(os.path.join(output_dir, "p_values.npy"), p_values)

    print(f"模型和检验值已保存到 {output_dir} 目录")


# -------------------- 核心功能函数 --------------------

def perform_regression_with_stats(ndvi_data, temp_data, precip_data, et_data):
    """
    执行多元回归分析，并计算R²、t值和p值
    :param ndvi_data: NDVI数据（三维数组，形状为(years, rows, cols)）
    :param temp_data: 温度数据（三维数组，形状为(years, rows, cols)）
    :param precip_data: 降水数据（三维数组，形状为(years, rows, cols)）
    :param et_data: 蒸散发数据（三维数组，形状为(years, rows, cols)）
    :return: 回归模型，回归系数（形状为(rows, cols, 3)），截距（形状为(rows, cols)），
             R²（形状为(rows, cols)），t值（形状为(rows, cols, 4)），p值（形状为(rows, cols, 4)）
    """
    # 获取图像的形状
    years, rows, cols = ndvi_data.shape

    # 将数据展平为二维数组 (years, rows*cols)
    ndvi_flat = ndvi_data.reshape(years, -1)
    temp_flat = temp_data.reshape(years, -1)
    precip_flat = precip_data.reshape(years, -1)
    et_flat = et_data.reshape(years, -1)

    # 初始化回归模型
    model = LinearRegression()

    # 逐像元回归
    coefficients = np.zeros((rows * cols, 3))  # 存储回归系数
    intercepts = np.zeros(rows * cols)  # 存储截距
    r2_values = np.zeros(rows * cols)  # 存储R²
    t_values = np.zeros((rows * cols, 4))  # 存储t值（包括截距）
    p_values = np.zeros((rows * cols, 4))  # 存储p值（包括截距）

    # 使用tqdm显示进度
    for i in tqdm(range(rows * cols), desc="逐像元回归", unit="像元"):
        # 提取当前像元的自变量和因变量
        X = np.column_stack((temp_flat[:, i], precip_flat[:, i], et_flat[:, i]))
        y = ndvi_flat[:, i]

        # 检查是否存在NaN或inf值
        valid = ~np.isnan(X).any(axis=1) & ~np.isnan(y) & ~np.isinf(X).any(axis=1) & ~np.isinf(y)
        X_valid = X[valid]
        y_valid = y[valid]

        # 如果有有效数据，进行回归
        if len(X_valid) > 0:
            model.fit(X_valid, y_valid)
            coefficients[i] = model.coef_
            intercepts[i] = model.intercept_

            # 计算R²
            y_pred = model.predict(X_valid)
            r2_values[i] = model.score(X_valid, y_valid)

            # 计算t值和p值
            n = len(X_valid)
            p = X_valid.shape[1] + 1  # 包括截距
            residuals = y_valid - y_pred
            mse = np.sum(residuals ** 2) / (n - p)
            X_with_intercept = np.column_stack((np.ones(n), X_valid))  # 添加截距项
            cov = np.linalg.inv(np.dot(X_with_intercept.T, X_with_intercept))
            se = np.sqrt(np.diag(mse * cov))
            t_values[i] = np.append(model.intercept_, model.coef_) / se  # 包括截距
            p_values[i] = 2 * (1 - stats.t.cdf(np.abs(t_values[i]), df=n - p))  # 计算p值
        else:
            # 如果没有有效数据，将系数、截距、R²、t值和p值设置为NaN
            coefficients[i] = np.nan
            intercepts[i] = np.nan
            r2_values[i] = np.nan
            t_values[i] = np.nan
            p_values[i] = np.nan

    # 将结果重新整形为图像形状
    coefficients = coefficients.reshape(rows, cols, 3)
    intercepts = intercepts.reshape(rows, cols)
    r2_values = r2_values.reshape(rows, cols)
    t_values = t_values.reshape(rows, cols, 4)  # 包括截距
    p_values = p_values.reshape(rows, cols, 4)  # 包括截距

    return model, coefficients, intercepts, r2_values, t_values, p_values


def predict_ndvi(coefficients, intercepts, temp_data, precip_data, et_data, template_file, output_dir="output"):
    """
    预测NDVI
    :param coefficients: 回归系数（形状为(rows, cols, 3)）
    :param intercepts: 截距（形状为(rows, cols)）
    :param temp_data: 温度数据（二维数组，形状为(rows, cols)）
    :param precip_data: 降水数据（二维数组，形状为(rows, cols)）
    :param et_data: 蒸散发数据（二维数组，形状为(rows, cols)）
    :param template_file: 模板文件路径（用于保存结果）
    :param output_dir: 输出目录（默认为"output"）
    :return: 预测的NDVI数据（二维数组，形状为(rows, cols)）
    """
    # 展平数据
    temp_flat = temp_data.reshape(-1)
    precip_flat = precip_data.reshape(-1)
    et_flat = et_data.reshape(-1)

    # 预测NDVI
    rows, cols = intercepts.shape
    ndvi_pred = np.zeros(rows * cols)

    # 使用tqdm显示进度
    for i in tqdm(range(rows * cols), desc="预测NDVI", unit="像元"):
        # 检查是否存在NaN或inf值
        if not np.isnan(coefficients.reshape(-1, 3)[i]).any() and \
                not np.isnan(temp_flat[i]) and \
                not np.isnan(precip_flat[i]) and \
                not np.isnan(et_flat[i]) and \
                not np.isinf(temp_flat[i]) and \
                not np.isinf(precip_flat[i]) and \
                not np.isinf(et_flat[i]):
            ndvi_pred[i] = (coefficients.reshape(-1, 3)[i][0] * temp_flat[i] +
                            coefficients.reshape(-1, 3)[i][1] * precip_flat[i] +
                            coefficients.reshape(-1, 3)[i][2] * et_flat[i] +
                            intercepts.reshape(-1)[i])
        else:
            ndvi_pred[i] = np.nan

    # 将预测结果重新整形为图像形状
    ndvi_pred = ndvi_pred.reshape(rows, cols)

    # 裁剪NDVI值到[0, 1]范围
    ndvi_pred = clip_ndvi(ndvi_pred)

    # 保存预测结果
    save_image(ndvi_pred, "ndvi_pred.tif", template_file, output_dir)

    return ndvi_pred


# -------------------- 主函数 --------------------

if __name__ == "__main__":
    # 定义数据目录
    workspace_dir = r'F:\test'
    ndvi_dir = r"F:\test\NDVI2" # NDVI数据目录, 范围0-1
    precip_dir = r'F:\test\pre'  # 降水数据目录
    temp_dir = r'F:\test\tmp'  # 平均气温数据目录
    et_dir = r'F:\test\pet'   # 蒸散发数据目录
    model_dir = "F:\test\model"  # 模型保存目录

    predict_dir = r'F:\test\ssp585'  # 用于预测的场景数据目录
    output_dir = r'F:\test\ssp585_pred'  # 输出目录


    predict_start_year = 2021  # 用于预测的起始年份
    predict_end_year = 2100  # 用于预测的结束年份


    # 用于获取情景的关键字
    keyword_tmp = "tmp"  # 命名规则tmp_2021这种格式
    keyword_pre = "pre"  # 同上
    keyword_pet = "pet"  # 同上


    # 读取NDVI数据（2000-2020年）
    ndvi_data = read_data_from_directory(ndvi_dir)
    ndvi_stack = np.stack(ndvi_data, axis=0)

    # 读取年均温数据（2000-2020年）
    temp_data = read_data_from_directory(temp_dir)
    temp_stack = np.stack(temp_data, axis=0)

    # 读取年总降水数据（2000-2020年）
    precip_data = read_data_from_directory(precip_dir)
    precip_stack = np.stack(precip_data, axis=0)

    # 读取年总蒸散发数据（2000-2020年）
    et_data = read_data_from_directory(et_dir)
    et_stack = np.stack(et_data, axis=0)

    # 执行多元回归分析并计算R²、t值和p值
    model, coefficients, intercepts, r2_values, t_values, p_values = perform_regression_with_stats(
        ndvi_stack, temp_stack, precip_stack, et_stack
    )

    # 保存模型和检验值
    save_model(model, coefficients, intercepts, r2_values, t_values, p_values, output_dir=model_dir)

    # 获取模板文件路径
    template_file = get_template_file(ndvi_dir)

    # 保存回归系数和截距
    save_image(coefficients[:, :, 0], "temp_coefficient.tif", template_file, model_dir)
    save_image(coefficients[:, :, 1], "precip_coefficient.tif", template_file, model_dir)
    save_image(coefficients[:, :, 2], "et_coefficient.tif", template_file, model_dir)
    save_image(intercepts, "intercept.tif", template_file, model_dir)

    # 保存R²、t值和p值
    save_image(r2_values, "r2_values.tif", template_file, model_dir)
    for i, var_name in enumerate(["intercept", "temp", "precip", "et"]):
        save_image(t_values[:, :, i], f"t_values_{var_name}.tif", template_file, model_dir)
        save_image(p_values[:, :, i], f"p_values_{var_name}.tif", template_file, model_dir)


    # 预测2021-2060年的NDVI值
    for year in range(predict_start_year, predict_end_year+1):
        print(f"正在预测 {year} 年的NDVI...")

        # 读取未来年份的温度、降水和蒸散发数据
        temp_future = read_data_from_directory(predict_dir, keyword=keyword_tmp + f"_{year}")
        precip_future = read_data_from_directory(predict_dir, keyword=keyword_pre + f"_{year}")
        et_future = read_data_from_directory(predict_dir, keyword=keyword_pet + f"_{year}")

        # 检查是否成功读取数据
        if len(temp_future) == 0 or len(precip_future) == 0 or len(et_future) == 0:
            print(f"警告：{year} 年的数据缺失，跳过该年份的预测。")
            continue

        # 提取数据（假设每年只有一个文件）
        temp_future = temp_future[0]
        precip_future = precip_future[0]
        et_future = et_future[0]

        # 预测NDVI
        ndvi_future_pred = predict_ndvi(coefficients, intercepts, temp_future, precip_future, et_future, template_file,
                                        output_dir)

        # 保存预测结果
        save_image(ndvi_future_pred, f"ndvi_pred_{year}.tif", template_file, output_dir)

    print(f"NDVI预测完成，结果已保存到 {output_dir} 目录")


